# Copyright (c) 2020 cgsdfc
#
# This software is released under the MIT License.
# https://opensource.org/licenses/MIT

from pathlib import Path

from astlib.asdl import parse as ParseAsdl, Module, VisitorBase

ASDL_INDEX_DATA_FILE = Path(__file__).with_suffix('.asdl')

assert ASDL_INDEX_DATA_FILE.exists()


class FieldVisitor(VisitorBase):
    def __init__(self):
        super().__init__()
        self.index = {}

    def visitModule(self, mod):
        for dfn in mod.dfns:
            self.visit(dfn)

    def visitType(self, type):
        self.visit(type.value, str(type.name))

    def visitSum(self, sum, name):
        for t in sum.types:
            self.visit(t, name)

    def visitConstructor(self, cons, name):
        key = str(cons.name)
        for f in cons.fields:
            self.visit(f, key)

    def visitField(self, field, name):
        key = (name, field.name)
        assert key not in self.index
        self.index[key] = field

    def visitProduct(self, prod, name):
        for f in prod.fields:
            self.visit(f, name)


class AsdlIndex:
    def __init__(self):
        self._mod = ParseAsdl(filename=str(ASDL_INDEX_DATA_FILE))
        vis = FieldVisitor()
        vis.visit(self._mod)
        self._field_index = vis.index

    def FindField(self, name, field_name):
        return self._field_index[name, field_name]


if __name__ == "__main__":
    asdl_index = AsdlIndex()
    print(asdl_index._field_index)
